#include <iostream>
#include <omp.h>
#include <libpmem.h>
#include <unistd.h>
#include <cassert>
#include <vector>
#include <random>
#include <algorithm>

#include "common.h"
#include "access.h"

using std::cout;
using std::endl;
using std::vector;

double run_test(char * buffer, bool read) {
    uint32_t count = BUFSIZE / ALIGN_SIZE;
    
    auto start = seconds();

    if(read) {
        uint64_t randnum = getRandom();
        uint64_t cur_offset = (randnum % BUFSIZE) ^ (ALIGN_SIZE - 1);

        for(uint32_t i = 0; i < count; i++) {
            cur_offset = *(uint64_t *)(buffer + cur_offset);
        }
        
        if(cur_offset < 256) printf("Avoiding optimize out");
    } else {
        rand_write_lat(buffer, count);
    }

    auto end = seconds();
    return end - start;
}

void init_memory(char * buffer) {
    vector<uint64_t> vec(BUFSIZE / ALIGN_SIZE);

    for(uint64_t i = 0; i < (BUFSIZE / ALIGN_SIZE); i++) {
        vec[i] = i * ALIGN_SIZE;
    }

    std::default_random_engine rd(getRandom());
    std::shuffle(vec.begin(), vec.end(), rd);

    for(uint64_t i = 0; i < (BUFSIZE / ALIGN_SIZE); i++) {
        *(uint64_t *)(buffer + i * ALIGN_SIZE) = vec[i];
    }
}

void print_help(char * name) {
    cout << "Usage: " << name <<" options" << endl;
    cout << "\t -p: test on PMEM, if not specified, test on DRAM" << endl;
    cout << "\t -w: write latency, if not specified, test read latency" << endl;
    cout << "\t -r: randomly, if not specified, test sequentially" << endl;
    return ;
}

int main(int argc, char ** argv) {
    bool opt_read = true;
    bool opt_pmem = false;

    static const char * optstr = "pwh";
    char opt;
    while((opt = getopt(argc, argv, optstr)) != -1) {
        switch(opt) {
            case 'p':
                opt_pmem = true; 
                break;
            case 'w': 
                opt_read = false; 
                break;
            case 'h':
            case '?': print_help(argv[0]); exit(-1);
            default:  print_help(argv[0]); exit(-1);
        }
    }

    // print the options
    #ifdef DEBUG
        cout << "pmem: " << (opt_pmem ? 1 : 0) << endl;
        cout << "read: " << (opt_read ? 1 : 0) << endl;
    #endif


    // allocating memory
    char * buf0 = NULL, *buf;
    size_t mapped_len;
    if (opt_pmem) {
        int is_pmem;
        if(file_exist("/mnt/pmem/bandwidth.pool")) {
            if(BUFSIZE > file_size("/mnt/pmem/bandwidth.pool")) {
                cout << "\nThe existing file is too small to do the test" << endl;
                exit(1);
            }
            // open the existing pool file and mmap it into virtual memory
            if ((buf0 = (char *)pmem_map_file("/mnt/pmem/bandwidth.pool", 0, 0,
                        0666, &mapped_len, &is_pmem)) == NULL) {
                perror("pmem_map_file");
                exit(1);
            }
        } else {
            // create a new pool file and mmap it into virtual memory 
            if ((buf0 = (char *)pmem_map_file("/mnt/pmem/bandwidth.pool", BUFSIZE + ALIGN_SIZE, PMEM_FILE_CREATE,
                        0666, &mapped_len, &is_pmem)) == NULL) {
                perror("pmem_map_file");
                exit(1);
            }
            
            // an empty file must be initialized with contents before reading
            omp_set_num_threads(16);
            #pragma omp parallel for
                for(int64_t i = 0; i < (BUFSIZE + ALIGN_SIZE) / 64; i++) {
                    buf0[i * 64] = 0x5a;
                }
        }

        // whether the pool file locates on PMEM or normal file system? 
        if(is_pmem == 0) {
            cout << "This file is not on pmem" << endl;
            exit(1);
        }
        buf = buf0;
    } else {
        buf0 = (char *) malloc(BUFSIZE + ALIGN_SIZE);
        uint64_t tmp = (uint64_t)buf0;
        uint64_t remaining = tmp & (ALIGN_SIZE - 1);
        buf = (char *) (remaining == 0 ? tmp : tmp + ALIGN_SIZE - remaining);
    }

    init_memory(buf);

    double elaspe_time = run_test(buf, opt_read);

    printf("%4.1lfns\n", elaspe_time / (BUFSIZE / ALIGN_SIZE) * 1000000000);

    // free memory
    if (opt_pmem)
        pmem_unmap(buf0, mapped_len);
    else
        free(buf0);
    
    return 0;
}